path= "/home/ma-user/work/Garbage_Data"
import mindspore.dataset as ds
import mindspore.dataset.vision as vision
 
def create_dataset(path, batch_size=50, train=True, image_size=224):
    #dataset = ds.ImageFolderDataset(path, num_parallel_workers=8, class_indexing={"croissants": 0, "dog": 1})
    dataset = ds.ImageFolderDataset(path, num_parallel_workers=8, class_indexing={"cardboard": 0, "glass": 1, "metal": 2, "paper": 3, "plastic": 4, "trash": 5})
 
    # 图像增强操作
    mean = [0.485 * 255, 0.456 * 255, 0.406 * 255]
    std = [0.229 * 255, 0.224 * 255, 0.225 * 255]
    if train:
        trans = [
            vision.RandomCropDecodeResize(image_size, scale=(0.08, 1.0), ratio=(0.75, 1.333)),
            vision.RandomHorizontalFlip(prob=0.5),
            vision.Normalize(mean=mean, std=std),
            vision.HWC2CHW()
        ]
    else:
        trans = [
            vision.Decode(),
            vision.Resize(256),
            vision.CenterCrop(image_size),
            vision.Normalize(mean=mean, std=std),
            vision.HWC2CHW()
        ]
 
    dataset = dataset.map(operations=trans, input_columns="image", num_parallel_workers=8)
    # 设置batch_size的大小，若最后一次抓取的样本数小于batch_size，则丢弃
    dataset = dataset.batch(batch_size, drop_remainder=True)
    return dataset
 
# 加载训练数据集和验证数据集用于后续的模型训练和验证。
# 加载训练数据集
train_path = "Garbage_Data/train"
dataset_train = create_dataset(train_path, train=True)
 
# 加载验证数据集
val_path = "Garbage_Data/val"
dataset_val = create_dataset(val_path, train=False)
 
 

#下载预训练模型文件 
from mindvision.dataset import DownLoad
models_url = "https://download.mindspore.cn/vision/classification/mobilenet_v2_1.0_224.ckpt"
dl = DownLoad()
# 下载模型文件
dl.download_url(models_url)
 


# 本章使用MobileNet V2的预训练模型进行微调，通过删除MobileNet V2预训练模型中最后一个用于分类的1x1的卷积层的参数，使用狗与牛角包数据集对模型进行重新训练以更新模型参数。 
import mindspore.nn as nn
import mindspore as ms
 
from mindvision.classification.models import mobilenet_v2
from mindvision.engine.loss import CrossEntropySmooth
 
# 创建模型,其中目标分类数为6，图像输入大小为(224,224)
network = mobilenet_v2(num_classes=6, resize=224)
# 模型参数存入到param_dict
param_dict = ms.load_checkpoint("./mobilenet_v2_1.0_224.ckpt")
# 获取mobilenet_v2网络最后一个卷积层的参数名
filter_list = [x.name for x in network.head.classifier.get_parameters()]
 
# 删除预训练模型的最后一个卷积层
def filter_ckpt_parameter(origin_dict, param_filter):
    for key in list(origin_dict.keys()):
        for name in param_filter:
            if name in key:
                print("Delete parameter from checkpoint: ", key)
                del origin_dict[key]
                break
 
filter_ckpt_parameter(param_dict, filter_list)
# 加载预训练模型参数作为网络初始化权重
ms.load_param_into_net(network, param_dict)
# 定义优化器
network_opt = nn.Momentum(params=network.trainable_params(), learning_rate=0.01, momentum=0.9)
# 定义损失函数
network_loss = CrossEntropySmooth(sparse=True, reduction="mean", smooth_factor=0.1, classes_num=2) 
# 定义评价指标
metrics = {"Accuracy": nn.Accuracy()}
# 初始化模型
model = ms.Model(network, loss_fn=network_loss, optimizer=network_opt, metrics=metrics)
 



# ### 模型训练与评估
# 训练并评估网络，使用MindSpore Vision中的`mindvision.engine.callback.ValAccMonitor`接口打印训练的损失值和评估精度，且在训练完成后，保存评估精度最高的CKPT文件`best.ckpt`在当前目录下。
from mindvision.engine.callback import ValAccMonitor
import mindspore as ms
 
num_epochs = 1
# 模型训练与验证，训练完成后保存验证精度最高的ckpt文件（best.ckpt）到当前目录下
model.train(num_epochs,
            dataset_train,
            callbacks=[ValAccMonitor(model, dataset_val, num_epochs), ms.TimeMonitor()])
 



# ### 可视化模型推理预测
# 定义 `visualize_model` 函数，使用上述验证精度最高的模型对输入图像进行预测，并将预测结果可视化。
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image
 
import mindspore as ms
 
def visualize_model(path):
    image = Image.open(path).convert("RGB")
    image = image.resize((224, 224))
    plt.imshow(image)
 
    # 归一化处理
    mean = np.array([0.485 * 255, 0.456 * 255, 0.406 * 255])
    std = np.array([0.229 * 255, 0.224 * 255, 0.225 * 255])
    image = np.array(image)
    image = (image - mean) / std
    image = image.astype(np.float32)
 
    # 图像通道由(h, w, c)转换为(c, h, w)
    image = np.transpose(image, (2, 0, 1))
 
    # 扩展数据维数为(1, c, h, w)
    image = np.expand_dims(image, axis=0)
 
    # 定义并加载网络
    net = mobilenet_v2(num_classes=6, resize=224)
    param_dict = ms.load_checkpoint("./best.ckpt")
    ms.load_param_into_net(net, param_dict)
    model = ms.Model(net)
 
    # 模型预测
    pre = model.predict(ms.Tensor(image))
    result = np.argmax(pre)
 
    #class_name = {0: "Croissants", 1: "Dog"}
    class_name = {0:"cardboard", 1:"glass", 2:"metal", 3:"paper", 4: "plastic", 5: "trash"}
    plt.title(f"Predict: {class_name[result]}")
    return result
 
image1 = "Garbage_Data/infer/cardboard1.jpg"
plt.figure(figsize=(15, 7))
plt.subplot(1, 2, 1)
visualize_model(image1)
 
image2 = "Garbage_Data/infer/glass1.jpg"
plt.figure(figsize=(15, 7))
plt.subplot(1, 2, 2)
visualize_model(image2)
 
image3 = "Garbage_Data/infer/metal1.jpg"
plt.figure(figsize=(15, 7))
plt.subplot(1, 2, 2)
visualize_model(image3)
 
image4 = "Garbage_Data/infer/paper1.jpg"
plt.figure(figsize=(15, 7))
plt.subplot(1, 2, 2)
visualize_model(image4)
 
image5 = "Garbage_Data/infer/plastic1.jpg"
plt.figure(figsize=(15, 7))
plt.subplot(1, 2, 2)
visualize_model(image5)
 
image6 = "Garbage_Data/infer/trash1.jpg"
plt.figure(figsize=(15, 7))
plt.subplot(1, 2, 2)
visualize_model(image6)

plt.show()
 
 

# 模型导出
# 在模型训练完后，训练完成后的网络模型（即CKPT文件）转换为MindIR格式，用于后续手机侧的推理。通过`export`接口会在当前目录下会生成`mobilenet_v2_1.0_224.mindir`文件。
import mindspore as ms
 
# 定义并加载网络参数
net = mobilenet_v2(num_classes=6, resize=224)
param_dict = ms.load_checkpoint("best.ckpt")
ms.load_param_into_net(net, param_dict)
 
# 将模型由ckpt格式导出为MINDIR格式
input_np = np.random.uniform(0.0, 1.0, size=[1, 3, 224, 224]).astype(np.float32)
ms.export(net, ms.Tensor(input_np), file_name="mobilenet_v2_1.0_224", file_format="MINDIR")